import math
import window
import numpy
import scipy.signal as signal




#s: STFT, 
#f: @freq
#t: @ time
def stft(x,fs,  **kwargs):
# 
# Input:
#	x:data
#	win: window
# 
#	hop : hop size
#	nfft: FFT points	
#   fs: frame rate
#
# Output:
#	s: STFT[T,F]	
#	f-freq	
#	t-center time	
	nfft=128
	freq_range=0
	overlap_len=-1
	win=None
	for key in kwargs:
		v=kwargs[key]
		if key== "Window":
			win=v	
		if key== "FFTLength":
			nfft= int(v)
		elif key== "FrequencyRange":	
			if v=='centered':	   
				freq_range= 0
			elif v=='onesided':	   
				freq_range= 1
			elif v=='twosided':	   
				freq_range= 2			 
			else:
				raise ValueError('invalid freq range')
		elif key =="overlapLength":	   
			overlap_len=int(v)

#win:Spectral window, specified as a vector. If you do not specify the window or specify it as empty, 
#the function uses a Hann window of length 128. The length of Window must be greater than or equal to 2.
	
	if win is None:
		win=window.hann(128,'periodic')
	   
	xlen=len(x)
	wlen=len(win)
	
	#overlapLength:Number of overlapped samples, specified as a positive integer smaller than the length of Window. 
	if overlap_len>=wlen:
		raise ValueError('overlapLength too large')

    #If you omit OverlapLength or specify it as empty, it is set to the largest integer less than 75% of the window length, 
	#which is 96 samples for the default Hann window.		
	elif overlap_len<0:
		overlap_len=wlen*0.75
	    
	hop=wlen-overlap_len
	
	#win count L
	L=int((xlen-overlap_len)/hop) 
	
	if freq_range == 1:
	    points=int(nfft/2 + 1)
	else:
		points=nfft
	
	s = numpy.zeros([L,points],dtype= complex)

	xw=numpy.zeros(wlen,dtype= float)
	off=0
	# STFT
	for l in range(L):
		# apply window
		n=xlen-off
		if n>wlen:
			n=wlen

		for i in range(n):
			xw[i] = x[off+i]*win[i]
	
		#  calc FFT
		fft_=numpy.fft.fft(xw,nfft) 
		
		if freq_range==0:
				fft_ =numpy.fft.fftshift(fft_)

		if freq_range==1:		
			numpy.copyto(s[l,:],fft_[0:int(nfft/2+1)]	   )
		else:
			numpy.copyto(s[l,:],fft_[:])

		off=off + hop

	# center of each win
	t=[(wlen/2+i*hop)/fs for i in range(L)]
	if freq_range==0:
		#after fftshift
		f=[(i-nfft/2+1)*(fs/nfft) for i in range(nfft)]
	elif freq_range==1:

		f=[(i)*(fs/nfft) for i in range(1+int(nfft/2))]		
	else:
	    f=[i*fs/nfft for i in range (nfft)]
	
	return s,f,t

def print_a(s):
	for i in range(s.shape[0]):
		print(i, end=" ")
		for j in range(s.shape[1]):
		    print( '{:.4f}'.format(s[i][j]), end=",")
		print()
		

if __name__=='__main__':
	#import matplotlib.pyplot as plt
	from scipy.io import wavfile
	import sys
	import os

	fs = 16e3;


	if len(sys.argv)<2:
		fn = "stft.wav"
	else:
		fn = sys.argv[1]

	fs, data = wavfile.read(fn)

	x = data/ 32768


	base_filename = os.path.basename(fn).replace(".wav", "")
	numpy.save(base_filename+"-data",x)

	segmentDuration = 1;
	frameDuration = 0.025;
	hopDuration = 0.010;

	FFTLength = 512;
	numBands = 50;

	segmentSamples = round(segmentDuration * fs);
	frameSamples = round(frameDuration * fs);
	hopSamples = round(hopDuration * fs);
	overlapSamples = frameSamples - hopSamples;
#	for i in range(x.shape[0]):
#		print("%f"%x[i])
	L = frameSamples
	win=window.hann(L, 'periodic')

	numpy.save(base_filename+"-win",win)
#	print(win)
#	plt.figure(0)
	
#	plt.plot(x[0:200])
#	plt.show()		
#	s,f,t=stft(x,fs,Window=win,overlapLength=overlapSamples,FFTLength=FFTLength)
#	print('centered')
#	print_a(s)
#	print(f)
#	print(t)	
	
	s,f,t=stft(x,fs,Window=win,overlapLength=overlapSamples,FFTLength=FFTLength,FrequencyRange='onesided')

#	print('onesided')
#	print_a(s)
#	print(f)
#	print(t)	
	
#	s,f,t=stft(x,fs,Window=win,overlapLength=overlapSamples,FFTLength=FFTLength,FrequencyRange='twosided')
#	print('twosided')
#	print_a(s)
#	print(f)
#	print(t)
	numpy.save(base_filename+"-stft",s.T)